package com.gitee.feizns.explore.data_structure.union_find;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

/**
 * 399. 除法求值
 * @author feizns
 * @since 2020/3/11
 */
public class CalcEquation {

    static List<List<String>> q = Arrays.asList(Arrays.asList("a", "e"), Arrays.asList("b", "e"));
    static double[] values = { 4., 3. };
    static List<List<String>> queries = Arrays.asList(
            Arrays.asList("a", "b"),
            Arrays.asList("e", "e"),
            Arrays.asList("x", "x")
    );
//[["a","e"],["b","e"]]
//[4.0,3.0]
//[["a","b"],["e","e"],["x","x"]]
    public static void main(String[] args) {
        System.out.println(Arrays.toString(calcEquation(q, values, queries)));
    }

    public static double[] calcEquation(List<List<String>> equations, double[] values, List<List<String>> queries) {
        Uf uf = new Uf(equations.stream().flatMap(List::stream));
        for (int i = 0; i < equations.size(); i++) {
            List<String> equation = equations.get(i);
            String a = equation.get(0);
            String b = equation.get(1);
            uf.union(b, a, values[i]);
        }

        double[] ret = new double[queries.size()];
        for (int i = 0; i < ret.length; i++) {
            List<String> que = queries.get(i);
            String a = que.get(0);
            String b = que.get(1);
            ret[i] = uf.multiple(a, b);
        }
        return ret;
    }

    static class Uf {
        public Uf(Stream<String> set) {
            set.forEach(s -> PARENT.put(s, s));
            for (String s : PARENT.keySet())
                VALUES.put(s, 1.);
        }
        Map<String, String> PARENT = new HashMap<>();
        Map<String, Double> VALUES = new HashMap<>();       //当前节点几倍能走到他父亲.
        void union(String a, String b, double val) {
            String ar = find(a);
            String br = find(b);
            if ( ar != null && br != null && !ar.equals(br) ) {
                PARENT.put(ar, br);
                VALUES.put(ar, (val * multiple(b)) / multiple(a));
            }
        }
        String find(String a) {
            String tmp = PARENT.get(a);
            if ( tmp != null ) {
                if ( !tmp.equals(a) ) {
                    String p = find(tmp);
                    VALUES.put(a, multiple(a));
                    PARENT.put(a, p);
                    return p;
                }
                return tmp;
            }
            return null;
        }
        //几倍到达根节点
        Double multiple(String a) {
            String s = PARENT.get(a);
            if ( a.equals(s) )
                return VALUES.get(a);
            return VALUES.get(a) * multiple(s);
        }
        //a是b的几倍
        Double multiple(String a, String b) {
            if ( isConnected(a, b) )
                return multiple(b) / multiple(a);
            return -1.;
        }
        boolean isConnected(String a, String b) {
            String ar = find(a);
            String br = find(b);
            return ar != null && br != null && ar.equals(br);
        }
    }

}
